# -*- coding: utf-8 -*-
"""
反演三条支路流量（Branch1/2/3）
Author: your_name
Date: 2025-07-11
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from scipy.optimize import nnls

# -------------------------------------------------
# 1️⃣ 读 Excel 第三张表（索引 2）
# -------------------------------------------------
FILE_PATH = '附件(Attachment).xlsx'
df_raw = pd.read_excel(FILE_PATH, sheet_name=2)

# 假设只有两列：时刻 & 主路4车流量
df_raw.columns = ['time_str', 'A3']

# 把字符串时间转成整数 t（07:00→0，07:02→1 … 08:58→59）
start_time = pd.to_datetime('07:00')
df_raw['t'] = ((pd.to_datetime(df_raw['time_str']) - start_time)
               .dt.total_seconds() / 120).astype(int)

# 排序并整理
df = df_raw[['t', 'A3']].sort_values('t').reset_index(drop=True)

# -------------------------------------------------
# 2️⃣ 构造设计矩阵 X 与观测向量 y
#    y = A3(t) = f1(t-1) + f2(t-1) + f3(t)
# -------------------------------------------------
T = 60
y = df['A3'].values          # shape (60,)
X = np.zeros((T, 3*T))       # shape (60,180)

for t in range(T):
    # 支路1/2 延迟 1 个时间单位
    if t - 1 >= 0:
        X[t, t - 1] = 1               # f1(t-1)
        X[t, T + t - 1] = 1           # f2(t-1)
    # 支路3 无延迟
    X[t, 2*T + t] = 1                 # f3(t)

# -------------------------------------------------
# 3️⃣ 非负最小二乘求解
# -------------------------------------------------
x_hat, _ = nnls(X, y)
f1 = x_hat[:T]
f2 = x_hat[T:2*T]
f3 = x_hat[2*T:]

# -------------------------------------------------
# 4️⃣ 重构主路4并计算 RMSE
# -------------------------------------------------
A3_pred = np.zeros(T)
for t in range(T):
    if t - 1 >= 0:
        A3_pred[t] += f1[t-1] + f2[t-1]
    A3_pred[t] += f3[t]

rmse = np.sqrt(np.mean((A3_pred - y)**2))
print(f'RMSE = {rmse:.3f}')

# -------------------------------------------------
# 5️⃣ 可视化
# -------------------------------------------------
plt.figure(figsize=(13,4))
sns.lineplot(x=df['t'], y=df['A3'], marker='o', label='Observed Main Road 4')
sns.lineplot(x=df['t'], y=A3_pred, label='Reconstructed Main Road 4')
plt.title('Observed vs Reconstructed Traffic Flow on Main Road 4')
plt.xlabel('t (2-min intervals since 07:00)')
plt.ylabel('Traffic Flow')
plt.legend(); plt.grid(True); plt.tight_layout()
plt.show()

# 三条支路单独作图
fig, ax = plt.subplots(3,1, figsize=(13,8), sharex=True)
ax[0].step(np.arange(T), f1, where='post'); ax[0].set_title('Branch 1 Flow')
ax[1].step(np.arange(T), f2, where='post'); ax[1].set_title('Branch 2 Flow')
ax[2].step(np.arange(T), f3, where='post'); ax[2].set_title('Branch 3 Flow')
ax[2].set_xlabel('t (2-min intervals since 07:00)')
for a in ax: a.grid(True)
plt.tight_layout(); plt.show()

# -------------------------------------------------
# 6️⃣ 输出表 3.2
# -------------------------------------------------
t730 = 15   # 7:30
t830 = 45   # 8:30
table32 = pd.DataFrame({
    'Time': ['07:30', '08:30'],
    'Branch1': [f1[t730-1]],
    'Branch2': [f2[t730-1]],
    'Branch3': [f3[t730]],
    'Branch1_8:30': [f1[t830-1]],
    'Branch2_8:30': [f2[t830-1]],
    'Branch3_8:30': [f3[t830]]
})
print('\n表 3.2  7:30 与 8:30 各支路流量估算')
print(table32.to_string(index=False))